/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "reachability_map.h"

namespace fe {
ReachabilityMap::ReachabilityMap(ge::ComputeGraph &graph) : size_(graph.GetDirectNode().size()) {
  bit_vectors_.reserve(size_);
  int64_t index_loop = 0;
  for (const auto &node : graph.GetDirectNode()) {
    name_to_index_[node->GetName()] = index_loop;
    indices_[index_loop] = bit_vectors_.size();
    bit_vectors_.emplace_back(size_);
    index_loop++;
  }
};

ReachabilityMap::~ReachabilityMap() {}

void ReachabilityMap::FastSetReachabilityToUnion(const std::vector<ge::NodePtr> &inputs, const ge::NodePtr &node) {
  SetReachabilityToUnionHelper(inputs, node, &GetBitVector(node));
}

void ReachabilityMap::SetReachabilityToUnionHelper(const std::vector<ge::NodePtr> &inputs, const ge::NodePtr &node,
                                                   BitVector *bit_vector) {
  // If instruction is part of inputs, don't reset the bit_vector.
  if (std::find(inputs.begin(), inputs.end(), node) == inputs.end()) {
    bit_vector->SetToZero();
  }
  bit_vector->Set(GetIndex(node));
  for (const ge::NodePtr &input : inputs) {
    if (input != node) {
      bit_vector->OrWith(GetBitVector(input));
    }
  }
}

bool ReachabilityMap::IsReachable(const ge::NodePtr &a, const ge::NodePtr &b) const {
  return GetBitVector(b).Get(GetIndex(a));
}

std::shared_ptr<ReachabilityMap> ReachabilityMap::Build(ge::ComputeGraph &graph) {
  std::shared_ptr<ReachabilityMap> result = std::shared_ptr<ReachabilityMap>(new ReachabilityMap(graph));
  FE_CHECK((result == nullptr), FE_LOGE("Reachability Map alloc failed"), return nullptr);
  std::vector<ge::NodePtr> inputs;

  for (auto &node : graph.GetDirectNode()) {
    inputs.clear();
    for (auto &input_node : node->GetInAllNodes()) {
      inputs.push_back(input_node);
    }
    result->FastSetReachabilityToUnion(inputs, node);
  }
  return result;
}

void ReachabilityMap::Update(ge::ComputeGraph &graph, vector<ge::NodePtr> &fusion_nodes) {
  BitVector new_bit_vector(graph.GetDirectNode().size());
  new_bit_vector.SetToZero();
  for (size_t i = 0; i < fusion_nodes.size(); i++) {
    new_bit_vector.OrWith(GetBitVector(fusion_nodes[i]));
  }
  for (auto &node : graph.GetDirectNode()) {
    bool is_connected_to_fusion = false;
    for (size_t i = 0; i < fusion_nodes.size(); i++) {
      if (GetBitVector(node).Get(GetIndex(fusion_nodes[i]))) {
        is_connected_to_fusion = true;
        break;
      }
    }
    if (is_connected_to_fusion) {
      GetBitVector(node).OrWith(new_bit_vector);
    }
  }
}
}
